# 导入torch工具包
import torch

# 指定source源
source = 'huggingface/pytorch-transformers'

# 指定part,代表加载的模块部分
part = 'tokenizer'

# 指定加载的模型, 使用自己上传的模型
model_name = 'ItcastAI/bert_finetunning_test'

# 首先导入映射器
tokenizer = torch.hub.load('huggingface/pytorch-transformers', 'tokenizer', model_name)

# 然后导入模型
model = torch.hub.load('huggingface/pytorch-transformers', 'modelForSequenceClassification', model_name)

# 进行预训练模型的编码
index = tokenizer.encode("Talk is cheap", "Please show me your code!")

# 指定BERT模型的句子间隔符, 102
mark = 102
# 找到第一个102的索引, 代表是两个句子的中间间隔符位置
k = index.index(mark)

# 构造一个列表结构, 来标志句子一的部分和句子二的部分
segments_ids = [0]*(k + 1) + [1]*(len(index) - k - 1)

# 将映射后的数字化张量和位置列表进行torch.tensor的封装
tokens_tensor = torch.tensor([index])
segments_tensor = torch.tensor([segments_ids])

# 使用预训练模型进行评估, 评估的时候注意模型参数不变化
with torch.no_grad():
    result = model(tokens_tensor, token_type_ids=segments_tensor)
    print(result)
    print(result[0].shape)



